from lxml import etree
from tqdm import tqdm as tqdm
import os, shutil
import os.path as osp
import pandas as pd


def recursive_parse_xml_to_dict(xml):
    if not xml:
        return {xml.tag: xml.text}
    result = {}
    for child in xml:
        child_result = recursive_parse_xml_to_dict(child)
        if child.tag != 'object':
            result[child.tag] = child_result[child.tag]
        else:
            if child.tag not in result:
                result[child.tag] = []
            result[child.tag].append(child_result[child.tag])
    return {xml.tag: result}


from xml.dom.minidom import Document


class VOC_XML():
    def __init__(self, image_path, image_width, image_height):
        self.image_path = image_path
        self.image_width = image_width
        self.image_height = image_height
        self.init_base_info()

    def init_base_info(self):
        self.doc = Document()  # 创建DOM文档对象
        self.root_element = self.doc.createElement('annotation')  # 创建根元素
        self.root_element.setAttribute('verified', 'no')
        self.doc.appendChild(self.root_element)
        # 创建folder节点
        folder = self.doc.createElement('folder')
        folder_text = self.doc.createTextNode('JPEGImages')
        folder.appendChild(folder_text)
        self.root_element.appendChild(folder)
        # 创建filename节点
        filename = self.doc.createElement('filename')
        filename_text = self.doc.createTextNode('000001')
        filename.appendChild(filename_text)
        self.root_element.appendChild(filename)
        # 创建path节点
        path = self.doc.createElement('path')
        path_text = self.doc.createTextNode(self.image_path)
        path.appendChild(path_text)
        self.root_element.appendChild(path)
        # 创建source节点
        source = self.doc.createElement('source')
        # 创建source节点的子节点database
        database = self.doc.createElement('database')
        database_text = self.doc.createTextNode('Unknown')
        database.appendChild(database_text)
        source.appendChild(database)
        self.root_element.appendChild(source)
        # 创建图片的size节点
        size = self.doc.createElement('size')
        # 创建size节点的子节点width节点
        width = self.doc.createElement('width')
        width_text = self.doc.createTextNode(str(int(self.image_width)))
        width.appendChild(width_text)
        size.appendChild(width)
        # 创建size节点的子节点width节点
        height = self.doc.createElement('height')
        height_text = self.doc.createTextNode(str(int(self.image_height)))
        height.appendChild(height_text)
        size.appendChild(height)
        # 创建size节点的子节点depth节点
        depth = self.doc.createElement('depth')
        depth_text = self.doc.createTextNode("3")
        depth.appendChild(depth_text)
        size.appendChild(depth)
        self.root_element.appendChild(size)
        # 创建segmented节点
        segmented = self.doc.createElement('segmented')
        segmented_text = self.doc.createTextNode('0')
        segmented.appendChild(segmented_text)
        self.root_element.appendChild(segmented)

    def add_object_element(self, cls, box, x_axis_first=True, _score=None):
        """
        :param cls: 类别名
        :param box: 坐标值
        :param x_axis_first:  True: [xmin,ymin,xmax,ymax]    False:[ymin,xmin,ymax,xmax]
        :param _score: 是否有，主要自动标注时可能出现
        :return:
        """
        if x_axis_first:
            _xmin, _ymin, _xmax, _ymax, = box
        else:
            _ymin, _xmin, _ymax, _xmax = box
        object_element = self.doc.createElement('object')
        # 创建object节点的子节点name节点
        name = self.doc.createElement('name')
        name_text = self.doc.createTextNode(cls)
        name.appendChild(name_text)
        object_element.appendChild(name)
        if not _score is None:
            # 创建object节点的子节点score节点
            score = self.doc.createElement('score')
            score_text = self.doc.createTextNode(str(_score))
            score.appendChild(score_text)
            object_element.appendChild(score)
        # 创建object节点的子节点pose节点
        pose = self.doc.createElement('pose')
        pose_text = self.doc.createTextNode('Unspecified')
        pose.appendChild(pose_text)
        object_element.appendChild(pose)
        # 创建object节点的子节点truncated节点
        truncated = self.doc.createElement('truncated')
        truncated_text = self.doc.createTextNode('0')
        truncated.appendChild(truncated_text)
        object_element.appendChild(truncated)
        # 创建object节点的子节点Difficult节点
        Difficult = self.doc.createElement('difficult')
        Difficult_text = self.doc.createTextNode('0')
        Difficult.appendChild(Difficult_text)
        object_element.appendChild(Difficult)
        # 创建object节点的子节点bndbox节点
        bndbox_element = self.doc.createElement('bndbox')
        # 创建bndbox的子节点xmin

        xmin = self.doc.createElement('xmin')
        xmin_text = self.doc.createTextNode(str(int(float(_xmin))))
        xmin.appendChild(xmin_text)
        bndbox_element.appendChild(xmin)
        # 创建bndbox的子节点ymin
        ymin = self.doc.createElement('ymin')
        ymin_text = self.doc.createTextNode(str(int(float(_ymin))))
        ymin.appendChild(ymin_text)
        bndbox_element.appendChild(ymin)
        # 创建bndbox的子节点xmax
        xmax = self.doc.createElement('xmax')
        xmax_text = self.doc.createTextNode(str(int(float(_xmax))))
        xmax.appendChild(xmax_text)
        bndbox_element.appendChild(xmax)
        # 创建bndbox的子节点ymax
        ymax = self.doc.createElement('ymax')
        ymax_text = self.doc.createTextNode(str(int(float(_ymax))))
        ymax.appendChild(ymax_text)
        bndbox_element.appendChild(ymax)
        bndbox_element.appendChild(ymax)
        object_element.appendChild(bndbox_element)
        self.root_element.appendChild(object_element)

    def write_file(self, xml_path):
        with open(xml_path, 'w', encoding='utf-8') as f:
            f.write(self.doc.toprettyxml(indent='  '))


def check_change():
    "用于测试转换前后各类别的数目是否对得上"
    txt_path = r"/Users/edy/Desktop/My_Projects/code_aculat/xml_each_category_num.txt"
    out_txt_path = r"/Users/edy/Desktop/My_Projects/code_aculat/xml_each_category_num_out.txt"

    with open(txt_path, 'r') as f:
        txt_list = f.read().splitlines()

    with open(out_txt_path, 'r') as f:
        out_txt_list = f.read().splitlines()

    txt_list_dict = {}
    out_txt_list_dict = {}

    for rec in txt_list:
        if len(rec.split(' ')) > 2:
            name_num = rec.split("%s " % rec.split(' ')[0])[1]
            name = name_num.split(" %s" % name_num.split(' ')[-1])[0]
            num = int(name_num.split(' ')[-1])
            txt_list_dict[name] = num
        else:
            name, num = rec.split(' ')
            num = int(num)
            txt_list_dict[name] = num

    for rec in out_txt_list:
        if len(rec.split(' ')) > 2:
            name_num = rec.split("%s " % rec.split(' ')[0])[1]
            name = name_num.split(" %s" % name_num.split(' ')[-1])[0]
            num = int(name_num.split(' ')[-1])
            out_txt_list_dict[name] = num
        else:
            name, num = rec.split(' ')
            num = int(num)
            out_txt_list_dict[name] = num

    for ind in txt_list_dict.keys():
        if ind not in out_txt_list_dict:
            print('%s  not in out ,num is %d' % (ind, txt_list_dict[ind]))
            continue
        if out_txt_list_dict[ind] != txt_list_dict[ind]:
            print('not equal', ind, out_txt_list_dict[ind], txt_list_dict[ind])


def change_name_accord_fix():
    xml_dir = r"/Users/edy/Data/RC/dp_add/Annotations"
    xml_out_dir = osp.join(osp.dirname(xml_dir), 'out_ann')
    os.makedirs(xml_out_dir, exist_ok=True)

    # raw_to_new_category = {}
    # manual_add_replace = {
    #     "东鹏加気能量型营养素饮料罐装355毫升": "DPS332F001S01T1W24P2355 东鹏加気能量型营养素饮料罐装355毫升",
    #     "东鹏特饮维生素功能饮料瓶装250毫升": "DPS002F002S01T1W6P1250 东鹏特饮维生素功能饮料瓶装250毫升",
    #     "东鹏特饮维生素功能饮料瓶装500毫升": "DPS002F002S01T1W6P1500 东鹏特饮维生素功能饮料瓶装500毫升",
    #     "东鹏特饮维生素功能饮料罐装250毫升": "DPS002F002S01T1W6P2250 东鹏特饮维生素功能饮料罐装250毫升"
    #     # "DPS002F002S01T1W6P1250 东鹏特饮维生素功能饮料罐装250毫升": "DPS002F002S01T1W6P2250 东鹏特饮维生素功能饮料罐装250毫升"
    # }
    # raw_category_list = ['DP1 C250ml-金罐',
    #                      'DP10 P400ml-陈皮特饮',
    #                      'DP11 C355ml-东鹏加气',
    #                      'DP12 C310ml-油柑柠檬茶',
    #                      'DP18 WATER380ml-水',
    #                      'DP2 P250ml-金瓶',
    #                      'DP20 C335ml-东鹏0糖',
    #                      'DP21 P330ml-大咖',
    #                      'DP22 P500ml-大咖',
    #                      'DP5 P500ml-金瓶',
    #                      'K0ANVCF00010001000 Never Coffee-咖啡饮料-拿铁咖啡-270ml',
    #                      'K0ANVCF00020001000 Never Coffee-咖啡饮料-燕麦咖啡-270ml',
    #                      'K0ANVCF00030001000 Never Coffee-咖啡饮料-美式咖啡-270ml',
    #                      'K1AWXR00010001000 外星人-电解质水-荔枝海盐-电解质饮料-PET瓶-500ml',
    #                      'K1AWXR00020001000 外星人-电解质水-青柠-电解质饮料-PET瓶-500ml',
    #                      'K1AWXR00030001000 外星人-电解质水-西柚-电解质饮料-PET瓶-500ml',
    #                      'K1AWXR00040001000 外星人-能量饮料-经典-能量饮料-PET瓶-350ml',
    #                      'K1AWXR00040003000 外星人-能量饮料-经典-能量饮料-铝罐-330ml',
    #                      'K1AWXR00050001000 外星人-能量饮料-马黛茶-能量饮料-PET瓶-350ml',
    #                      'K1AWXR00050003000 外星人-能量饮料-马黛茶-能量饮料-铝罐-330ml',
    #                      'K1AWXR00070001000 外星人-能量饮料-乳酸菌-能量饮料-PET瓶-350ml',
    #                      'K1AWXR00070003000 外星人-能量饮料-乳酸菌-能量饮料-铝罐-330ml',
    #                      'K1AWXR00080001000 外星人-能量饮料-西柚-能量饮料-PET瓶-350ml',
    #                      'K1AWXR00080003000 外星人-能量饮料-西柚-能量饮料-铝罐-330ml',
    #                      'K1AWXR00090001000 外星人-能量饮料-彩虹独角兽-PET瓶-350ml',
    #                      'K1AYQSL00180001000 元气森林-元气森林-白桃-苏打气泡水-PET瓶-1250ml -福気',
    #                      'K1AYQSL00180003000 元气森林-元气森林-白桃-苏打气泡水-PET瓶-480ml',
    #                      'K1AYQSL00180007000 元气森林-元气森林-白桃-苏打气泡水-铝罐-200ml',
    #                      'K1AYQSL00190001000 元气森林-元气森林-卡曼橘-苏打气泡水-PET瓶-480ml',
    #                      'K1AYQSL00190005000 元气森林-元气森林-卡曼橘-苏打气泡水-铝罐-200ml',
    #                      'K1AYQSL00210001000 元气森林-元气森林-荔枝-苏打气泡水-PET瓶-480ml',
    #                      'K1AYQSL00220001000 元气森林-元气森林-海盐菠萝-苏打气泡水-PET-480ml',
    #                      'K1AYQSL00230001000 元气森林-元气森林-青瓜-苏打气泡水-PET瓶-480ml',
    #                      'K1AYQSL00240002000 元气森林-元气森林-乳酸菌-苏打气泡水-PET瓶-480ml',
    #                      'K1AYQSL00250002000 元气森林-元气森林-酸梅汁-苏打气泡水-PET瓶-480ml',
    #                      'K1AYQSL00260001000 元气森林-元气森林-樱花白葡萄-苏打气泡水-PET瓶-480ml',
    #                      'K1AYQSL00290001000 元气森林-燃茶-草莓茉莉-无糖乌龙茶-PET瓶-500ml',
    #                      'K1AYQSL00300001000 元气森林-燃茶-醇香-无糖乌龙茶-PET瓶-500ml',
    #                      'K1AYQSL00310001000 元气森林-燃茶-桃香-无糖乌龙茶-PET瓶-500ml',
    #                      'K1AYQSL00320001000 元气森林-燃茶-玄米-无糖乌龙茶-PET瓶-500ml',
    #                      'K1AYQSL00340002000 元气森林-元气森林乳茶-咖啡拿铁-乳茶-PET瓶-450ml',
    #                      'K1AYQSL00350002000 元气森林-元气森林乳茶-茉香奶绿-乳茶-PET瓶-450ml',
    #                      'K1AYQSL00360002000 元气森林-元气森林乳茶-浓香原味-乳茶-PET瓶-450ml',
    #                      'K1AYQSL00410001000 元气森林-元气森林-夏黑葡萄-苏打气泡水-PET瓶-480ml',
    #                      'K1AYQSL00420001000 元气森林-满分-西柚-果汁微气泡-PET瓶-380ml',
    #                      'K1AYQSL00430001000 元气森林-满分-夏黑葡萄-果汁微气泡-PET瓶-380ml',
    #                      'K1AYQSL00440001000 元气森林-满分-王林青苹果-果汁微气泡-PET瓶-380ml',
    #                      'K1AYQSL00450001000 元气森林-元气森林-王林青苹果-苏打气泡水-PET瓶-480ml',
    #                      'K1AYQSL00490001000 元气森林-燃茶-焙火-无糖乌龙茶-PET瓶-500ml',
    #                      'K1AYQSL00670001000 元气森林-元气森林-百香果-苏打气泡水-PET-480ml',
    #                      'K1AYQSL00680001000 元气森林-元气森林-石榴红树莓-苏打气泡水-PET-480ml',
    #                      'K1AYQSL00690001000 元气森林-纤茶-玉米须茶-PET-500ml',
    #                      'K1AYQSL00710001000 元气森林-元气森林-红香酥梨-苏打气泡水-PET-480ml',
    #                      'bottle',
    #                      'row',
    #                      '乐虎氨基酸维生素功能饮料瓶装380毫升',
    #                      '健力宝运动饮料橙蜜味瓶装560毫升',
    #                      '农夫山泉东方树叶茉莉花茶瓶装500毫升',
    #                      '农夫山泉尖叫多肽型瓶装550毫升',
    #                      '农夫山泉尖叫纤维型瓶装550毫升',
    #                      '农夫山泉水溶C100柠檬味复合果汁饮料瓶装445毫升',
    #                      '农夫山泉水溶C100西柚味复合果汁饮料瓶装445毫升',
    #                      '农夫山泉维他命水柑橘风味瓶装500毫升',
    #                      '农夫山泉茶π柚子绿茶瓶装500毫升',
    #                      '农夫山泉茶π柠檬红茶瓶装500毫升',
    #                      '农夫山泉茶π蜜桃乌龙茶瓶装500毫升',
    #                      '农夫山泉茶π西柚茉莉花茶瓶装500毫升',
    #                      '农夫山泉饮用天然水瓶装550毫升',
    #                      '加多宝凉茶植物饮料瓶装550毫升',
    #                      '加多宝凉茶植物饮料罐装310毫升',
    #                      '可口可乐瓶装300毫升',
    #                      '可口可乐瓶装600毫升',
    #                      '可口可乐罐装330毫升',
    #                      '可口可乐零度瓶装500毫升',
    #                      '和其正凉茶瓶装550毫升',
    #                      '娃哈哈AD钙奶瓶装450毫升',
    #                      '娃哈哈营养快线果汁牛奶瓶装500毫升',
    #                      '娃哈哈饮用纯净水瓶装596毫升',
    #                      '康师傅冰糖雪梨瓶装1000毫升',
    #                      '康师傅冰糖雪梨瓶装500毫升',
    #                      '康师傅冰红茶柠檬口味瓶装1000毫升',
    #                      '康师傅冰红茶柠檬口味瓶装500毫升',
    #                      '康师傅喝开水熟水饮用水塑料瓶550毫升',
    #                      '康师傅绿茶低糖瓶装500毫升',
    #                      '康师傅绿茶瓶装1000毫升',
    #                      '康师傅茉莉清茶瓶装1000毫升',
    #                      '康师傅茉莉清茶瓶装500毫升',
    #                      '康师傅茉莉蜜茶瓶装1000毫升',
    #                      '康师傅茉莉蜜茶瓶装500毫升',
    #                      '康师傅蜂蜜柚子茶瓶装500毫升',
    #                      '康师傅酸梅汤瓶装500毫升',
    #                      '怡宝葡萄假日瓶装440毫升',
    #                      '怡宝饮用纯净水瓶装1555毫升',
    #                      '怡宝饮用纯净水瓶装550毫升',
    #                      '景田饮用纯净水瓶装560毫升',
    #                      '王老吉凉茶植物饮料瓶装500毫升',
    #                      '王老吉凉茶植物饮料罐装310毫升',
    #                      '百事可乐瓶装1000毫升',
    #                      '百事可乐瓶装1250毫升',
    #                      '百事可乐瓶装600毫升',
    #                      '百事可乐罐装330毫升',
    #                      '百威淡色拉格啤酒550毫升',
    #                      '百岁山天然矿泉水瓶装570毫升',
    #                      '红牛罐装250毫升',
    #                      '统一冰红茶柠檬味瓶装500毫升',
    #                      '统一小茗同学冷泡茶青柠红茶瓶装480毫升',
    #                      '统一海之言海盐柠檬水饮料瓶装500毫升',
    #                      '统一绿茶低糖茉莉味瓶装500毫升',
    #                      '统一阿萨姆原味奶茶瓶装500毫升',
    #                      '维他柠檬茶瓶装500毫升',
    #                      '维他柠檬茶蜜桃茶瓶装500毫升',
    #                      '美年达橙味瓶装600毫升',
    #                      '美汁源果粒橙瓶装1250毫升',
    #                      '芬达橙味(扭扭瓶)瓶装500毫升',
    #                      '达利园青梅绿茶瓶装500毫升',
    #                      '达能脉动水蜜桃口味瓶装600毫升',
    #                      '达能脉动菠萝口味瓶装600毫升',
    #                      '达能脉动雪柚橘子口味瓶装600毫升',
    #                      '达能脉动青柠口味瓶装600毫升',
    #                      '雪碧清爽柠檬味瓶装600毫升',
    #                      '雪碧清爽柠檬味罐装330毫升',
    #                      '雪花8度勇闯天涯瓶装500毫升',
    #                      '雪花8度勇闯天涯罐装500毫升',
    #                      '青岛啤酒经典罐装500毫升', ]
    # new_category_list = ["DPS002F002S01T1W6P2250 东鹏特饮维生素功能饮料罐装250毫升",
    #                      "DPS056F090S01T1W30P1400 东鹏陈皮特饮瓶装400毫升",
    #                      "DPS332F001S01T1W24P2355 东鹏加気能量型营养素饮料罐装355毫升",
    #                      "DPS056F090S01T1W29P9310 东鹏由柑柠檬茶高罐装310毫升",
    #                      "DPS357F057S03T1W9P1380 饮用天然水",
    #                      "DPS002F002S01T1W6P1250 东鹏特饮维生素功能饮料瓶装250毫升",
    #                      "DPS002F001S03T1W6P2335 东鹏0糖无糖型能量饮料",
    #                      "DPS356F003S02T1W20P1330 东鹏大咖®咖啡饮料330毫升",
    #                      "DPS356F003S02T1W20P1500 东鹏大咖®咖啡饮料500毫升",
    #                      "DPS002F002S01T1W6P1500 东鹏特饮维生素功能饮料瓶装500毫升",
    #                      'K0ANVCF00010001000 Never Coffee-咖啡饮料-拿铁咖啡-270ml',
    #                      'K0ANVCF00020001000 Never Coffee-咖啡饮料-燕麦咖啡-270ml',
    #                      'K0ANVCF00030001000 Never Coffee-咖啡饮料-美式咖啡-270ml',
    #                      'K1AWXR00010001000 外星人-电解质水-荔枝海盐-电解质饮料-PET瓶-500ml',
    #                      'K1AWXR00020001000 外星人-电解质水-青柠-电解质饮料-PET瓶-500ml',
    #                      'K1AWXR00030001000 外星人-电解质水-西柚-电解质饮料-PET瓶-500ml',
    #                      'K1AWXR00040001000 外星人-能量饮料-经典-能量饮料-PET瓶-350ml',
    #                      'K1AWXR00040003000 外星人-能量饮料-经典-能量饮料-铝罐-330ml',
    #                      'K1AWXR00050001000 外星人-能量饮料-马黛茶-能量饮料-PET瓶-350ml',
    #                      'K1AWXR00050003000 外星人-能量饮料-马黛茶-能量饮料-铝罐-330ml',
    #                      'K1AWXR00070001000 外星人-能量饮料-乳酸菌-能量饮料-PET瓶-350ml',
    #                      'K1AWXR00070003000 外星人-能量饮料-乳酸菌-能量饮料-铝罐-330ml',
    #                      'K1AWXR00080001000 外星人-能量饮料-西柚-能量饮料-PET瓶-350ml',
    #                      'K1AWXR00080003000 外星人-能量饮料-西柚-能量饮料-铝罐-330ml',
    #                      'K1AWXR00090001000 外星人-能量饮料-彩虹独角兽-PET瓶-350ml',
    #                      'K1AYQSL00180001000 元气森林-元气森林-白桃-苏打气泡水-PET瓶-1250ml -福気',
    #                      'K1AYQSL00180003000 元气森林-元气森林-白桃-苏打气泡水-PET瓶-480ml',
    #                      'K1AYQSL00180007000 元气森林-元气森林-白桃-苏打气泡水-铝罐-200ml',
    #                      'K1AYQSL00190001000 元气森林-元气森林-卡曼橘-苏打气泡水-PET瓶-480ml',
    #                      'K1AYQSL00190005000 元气森林-元气森林-卡曼橘-苏打气泡水-铝罐-200ml',
    #                      'K1AYQSL00210001000 元气森林-元气森林-荔枝-苏打气泡水-PET瓶-480ml',
    #                      'K1AYQSL00220001000 元气森林-元气森林-海盐菠萝-苏打气泡水-PET-480ml',
    #                      'K1AYQSL00230001000 元气森林-元气森林-青瓜-苏打气泡水-PET瓶-480ml',
    #                      'K1AYQSL00240002000 元气森林-元气森林-乳酸菌-苏打气泡水-PET瓶-480ml',
    #                      'K1AYQSL00250002000 元气森林-元气森林-酸梅汁-苏打气泡水-PET瓶-480ml',
    #                      'K1AYQSL00260001000 元气森林-元气森林-樱花白葡萄-苏打气泡水-PET瓶-480ml',
    #                      'K1AYQSL00290001000 元气森林-燃茶-草莓茉莉-无糖乌龙茶-PET瓶-500ml',
    #                      'K1AYQSL00300001000 元气森林-燃茶-醇香-无糖乌龙茶-PET瓶-500ml',
    #                      'K1AYQSL00310001000 元气森林-燃茶-桃香-无糖乌龙茶-PET瓶-500ml',
    #                      'K1AYQSL00320001000 元气森林-燃茶-玄米-无糖乌龙茶-PET瓶-500ml',
    #                      'K1AYQSL00340002000 元气森林-元气森林乳茶-咖啡拿铁-乳茶-PET瓶-450ml',
    #                      'K1AYQSL00350002000 元气森林-元气森林乳茶-茉香奶绿-乳茶-PET瓶-450ml',
    #                      'K1AYQSL00360002000 元气森林-元气森林乳茶-浓香原味-乳茶-PET瓶-450ml',
    #                      'K1AYQSL00410001000 元气森林-元气森林-夏黑葡萄-苏打气泡水-PET瓶-480ml',
    #                      'K1AYQSL00420001000 元气森林-满分-西柚-果汁微气泡-PET瓶-380ml',
    #                      'K1AYQSL00430001000 元气森林-满分-夏黑葡萄-果汁微气泡-PET瓶-380ml',
    #                      'K1AYQSL00440001000 元气森林-满分-王林青苹果-果汁微气泡-PET瓶-380ml',
    #                      'K1AYQSL00450001000 元气森林-元气森林-王林青苹果-苏打气泡水-PET瓶-480ml',
    #                      'K1AYQSL00490001000 元气森林-燃茶-焙火-无糖乌龙茶-PET瓶-500ml',
    #                      'K1AYQSL00670001000 元气森林-元气森林-百香果-苏打气泡水-PET-480ml',
    #                      'K1AYQSL00680001000 元气森林-元气森林-石榴红树莓-苏打气泡水-PET-480ml',
    #                      'K1AYQSL00690001000 元气森林-纤茶-玉米须茶-PET-500ml',
    #                      'K1AYQSL00710001000 元气森林-元气森林-红香酥梨-苏打气泡水-PET-480ml',
    #                      'bottle',
    #                      'row',
    #                      'LHS061F002S01T1W6P1380 乐虎氨基酸维生素功能饮料瓶装380毫升',
    #                      'JLBS086F114S01T1W6P1560 健力宝运动饮料橙蜜味瓶装560毫升',
    #                      'NFSS025F001S03T1W17P1500 农夫山泉东方树叶茉莉花茶瓶装500毫升',
    #                      'NFSS029F058S01T1W6P1550 农夫山泉尖叫多肽型瓶装550毫升',
    #                      'NFSS029F005S01T1W6P1550 农夫山泉尖叫纤维型瓶装550毫升',
    #                      'NFSS030F036S01T1W18P1445 农夫山泉水溶C100柠檬味复合果汁饮料瓶装445毫升',
    #                      'NFSS030F240S01T1W7P1445 农夫山泉水溶C100西柚味复合果汁饮料瓶装445毫升',
    #                      'NFSS032F064S01T1W24P1500 农夫山泉维他命水柑橘风味瓶装500毫升',
    #                      'NFSS026F049S01T1W16P1500 农夫山泉茶π柚子绿茶瓶装500毫升',
    #                      'NFSS026F036S01T1W15P1500 农夫山泉茶π柠檬红茶瓶装500毫升',
    #                      'NFSS026F051S01T1W10P1500 农夫山泉茶π蜜桃乌龙茶瓶装500毫升',
    #                      'NFSS026F009S01T1W17P1500 农夫山泉茶π西柚茉莉花茶瓶装500毫升',
    #                      'NFSS024F057S03T1W9P1550 农夫山泉饮用天然水瓶装550毫升',
    #                      'JDBS096F001S01T1W30P1550 加多宝凉茶植物饮料瓶装550毫升',
    #                      'JDBS096F001S01T1W30P2310 加多宝凉茶植物饮料罐装310毫升',
    #                      'CCS003F001S01T1W2P1300 可口可乐瓶装300毫升',
    #                      'CCS003F001S01T1W2P1600 可口可乐瓶装600毫升',
    #                      'CCS003F001S01T1W2P2330 可口可乐罐装330毫升',
    #                      'CCS005F001S01T1W2P1500 可口可乐零度瓶装500毫升',
    #                      'HQZS108F001S01T1W30P1550 和其正凉茶瓶装550毫升',
    #                      'WHHS342F125S01T1W32P1450 娃哈哈AD钙奶瓶装450毫升',
    #                      'WHHS329F125S01T1W32P1500 娃哈哈营养快线果汁牛奶瓶装500毫升',
    #                      'WHHS133F001S03T1W9P1596 娃哈哈饮用纯净水瓶装596毫升',
    #                      'MKS138F100S01T1W7P11000 康师傅冰糖雪梨瓶装1000毫升',
    #                      'MKS138F100S01T1W7P1500 康师傅冰糖雪梨瓶装500毫升',
    #                      'MKS161F036S01T1W15P11000 康师傅冰红茶柠檬口味瓶装1000毫升',
    #                      'MKS161F036S01T1W15P1500 康师傅冰红茶柠檬口味瓶装500毫升',
    #                      'MKS166F001S03T1W9P1550 康师傅喝开水熟水饮用水塑料瓶550毫升',
    #                      'MKS164F001S02T1W16P1500 康师傅绿茶低糖瓶装500毫升',
    #                      'MKS164F001S01T1W16P11000 康师傅绿茶瓶装1000毫升',
    #                      'MKS164F147S01T1W30P11000 康师傅茉莉清茶瓶装1000毫升',
    #                      'MKS164F147S01T1W30P1500 康师傅茉莉清茶瓶装500毫升',
    #                      'MKS164F148S01T1W30P11000 康师傅茉莉蜜茶瓶装1000毫升',
    #                      'MKS164F148S01T1W30P1500 康师傅茉莉蜜茶瓶装500毫升',
    #                      'MKF339F023S01T1W30P1500 康师傅蜂蜜柚子茶瓶装500毫升',
    #                      'MKS164F112S01T1W7P1500 康师傅酸梅汤瓶装500毫升',
    #                      'YBS345F008S01T1W7P1440 怡宝葡萄假日瓶装440毫升',
    #                      'YBS176F057S03T1W9P11555 怡宝饮用纯净水瓶装1555毫升',
    #                      'YBS176F057S01T1W9P1550 怡宝饮用纯净水瓶装550毫升',
    #                      'JTS195F057S03T1W9P1560 景田饮用纯净水瓶装560毫升',
    #                      'WLJS235F001S01T1W46P1500 王老吉凉茶植物饮料瓶装500毫升',
    #                      'WLJS235F001S01T1W46P2310 王老吉凉茶植物饮料罐装310毫升',
    #                      'PCS010F012S01T1W2P11000 百事可乐瓶装1000毫升',
    #                      'PCS010F012S01T1W2P11250 百事可乐瓶装1250毫升',
    #                      'PCS010F012S01T1W2P11600 百事可乐瓶装600毫升',
    #                      'PCS010F012S01T1W2P2330 百事可乐罐装330毫升',
    #                      'BWS327F202S01T1W49P2550 百威淡色拉格啤酒550毫升',
    #                      'JTS194F057S03T1W9P1570 百岁山天然矿泉水瓶装570毫升',
    #                      'RBS012F001S03T1W6P2250 红牛罐装250毫升',
    #                      'TYS161F036S01T1W15P1500 统一冰红茶柠檬味瓶装500毫升',
    #                      'TYS252F177S01T1W15P1480 统一小茗同学冷泡茶青柠红茶瓶装480毫升',
    #                      'TYS254F227S01T1W7P1500 统一海之言海盐柠檬水饮料瓶装500毫升',
    #                      'TYS251F101S02T1W16P1500 统一绿茶低糖茉莉味瓶装500毫升',
    #                      'TYS014F001S01T1W4P1500 统一阿萨姆原味奶茶瓶装500毫升',
    #                      'WTS256F036S01T1W29P1500 维他柠檬茶瓶装500毫升',
    #                      'WTS256F051S01T1W30P1500 维他柠檬茶蜜桃茶瓶装500毫升',
    #                      'MNDS015F019S01T1W2P1600 美年达橙味瓶装600毫升',
    #                      'MZYS264F001S01T1W7P11250 美汁源果粒橙瓶装1250毫升',
    #                      'FDS023F019S03T1W2P1500 芬达橙味(扭扭瓶)瓶装500毫升',
    #                      'DLYS278F092S01T1W16P1500 达利园青梅绿茶瓶装500毫升',
    #                      'DNNS016F030S01T1W6P1600 达能脉动水蜜桃口味瓶装600毫升',
    #                      'DNNS016F034S01T1W6P1600 达能脉动菠萝口味瓶装600毫升',
    #                      'DNNS016F206S01T1W6P1600 达能脉动雪柚橘子口味瓶装600毫升',
    #                      'DNNS016F037S01T1W6P1600 达能脉动青柠口味瓶装600毫升',
    #                      'XBS017F036S01T1W2P1600 雪碧清爽柠檬味瓶装600毫升',
    #                      'XBS017F036S01T1W2P2330 雪碧清爽柠檬味罐装330毫升',
    #                      'XHS304F202S01T1W49P1500 雪花8度勇闯天涯瓶装500毫升',
    #                      'XHS304F202S01T1W49P2500 雪花8度勇闯天涯罐装500毫升',
    #                      'QDS355F202S01T1W49P2500 青岛啤酒经典罐装500毫升', ]
    #
    # for index in range(len(raw_category_list)):
    #     raw_to_new_category[raw_category_list[index]] = new_category_list[index]
    raw_to_new_category = {
        "DP1 C250ml-金罐":'东鹏特饮维生素功能饮料罐装250毫升',
        "DP2 P250ml-金瓶":'东鹏特饮维生素功能饮料瓶装250毫升',
        "DP5 P500ml-金瓶":'东鹏特饮维生素功能饮料瓶装500毫升',

    }
    manual_add_replace=[]
    if len(manual_add_replace):
        for key, value in manual_add_replace.items():
            raw_to_new_category[key] = value

    for xml_name in tqdm(filter(lambda x: x.endswith('.xml'), os.listdir(xml_dir))):

        xml_path = os.path.join(xml_dir, xml_name)
        with open(xml_path, 'rb') as f:
            xml_str = f.read()

        xml = etree.fromstring(xml_str)
        xml_dict = recursive_parse_xml_to_dict(xml)

        if 'path' not in xml_dict["annotation"]:  # 如果缺少这个判断，后续可能会报错
            xml_dict["annotation"]['path'] = xml_dict["annotation"]['filename']
        xml_object = VOC_XML(xml_dict["annotation"]["path"],  # 这是一个新的xml文件对象，所以后面是使用add_object_element
                             xml_dict["annotation"]["size"]["width"],
                             xml_dict["annotation"]["size"]["height"])

        if 'object' in xml_dict["annotation"]:
            for index, obj in enumerate(xml_dict["annotation"]["object"]):
                if obj['name'] in raw_to_new_category:
                    obj['name'] = raw_to_new_category[obj['name']]

                xml_object.add_object_element(obj['name'],
                                              [obj['bndbox']['xmin'], obj['bndbox']['ymin'],
                                               obj['bndbox']['xmax'],
                                               obj['bndbox']['ymax']])

            xml_object.write_file(os.path.join(xml_out_dir, xml_name))


def change_name_accord_csv():
    xml_dir = r"/Users/edy/Data/DB/label_data/Annotations"
    xml_out_dir = osp.join(osp.dirname(xml_dir), 'out_ann')
    os.makedirs(xml_out_dir, exist_ok=True)

    raw_to_new_category = {}  # 原来名称与现在名称的对应关系

    # 从csv获取对应关系
    csv_path = r"/Users/edy/Downloads/元气森林-贩冰冰-商品识别编码库 - 汇总.csv"

    data = pd.read_csv(csv_path)
    sku_code = data['sku_code']
    sku_name = data['sku_name']
    name_to_code_dict = {}
    for idx in range(len(sku_code)):
        name_to_code_dict[sku_name[idx]] = sku_code[idx]
        raw_to_new_category[sku_name[idx]] = "%s %s" % (sku_code[idx], sku_name[idx],)

    for xml_name in tqdm(filter(lambda x: x.endswith('.xml'), os.listdir(xml_dir))):

        xml_path = os.path.join(xml_dir, xml_name)
        with open(xml_path, 'rb') as f:
            xml_str = f.read()

        xml = etree.fromstring(xml_str)
        xml_dict = recursive_parse_xml_to_dict(xml)

        if 'path' not in xml_dict["annotation"]:  # 如果缺少这个判断，后续可能会报错
            xml_dict["annotation"]['path'] = xml_dict["annotation"]['filename']
        xml_object = VOC_XML(xml_dict["annotation"]["path"],  # 这是一个新的xml文件对象，所以后面是使用add_object_element
                             xml_dict["annotation"]["size"]["width"],
                             xml_dict["annotation"]["size"]["height"])

        if 'object' in xml_dict["annotation"]:
            for index, obj in enumerate(xml_dict["annotation"]["object"]):
                if obj['name'] in raw_to_new_category:
                    obj['name'] = raw_to_new_category[obj['name']]
                else:
                    aa = 1

                xml_object.add_object_element(obj['name'],
                                              [obj['bndbox']['xmin'], obj['bndbox']['ymin'],
                                               obj['bndbox']['xmax'],
                                               obj['bndbox']['ymax']])

            xml_object.write_file(os.path.join(xml_out_dir, xml_name))


if __name__ == '__main__':
    # change_name_accord_csv()
    # check_change()
    change_name_accord_fix()
